package simpledb.execution;

import simpledb.common.Type;
import simpledb.storage.*;

import java.util.ArrayList;
import java.util.HashMap;
import java.util.List;
import java.util.Map;

/**
 * Knows how to compute some aggregate over a set of StringFields.
 */
public class StringAggregator implements Aggregator {

    private static final long serialVersionUID = 1L;
    private int gbfield;
    private Type gbfieldtype;
    private int afield;
    private Op what;
    private Map<Field,Integer> groups;
    private int nogroup;

    /**
     * Aggregate constructor
     * @param gbfield the 0-based index of the group-by field in the tuple, or NO_GROUPING if there is no grouping
     * @param gbfieldtype the type of the group by field (e.g., Type.INT_TYPE), or null if there is no grouping
     * @param afield the 0-based index of the aggregate field in the tuple
     * @param what aggregation operator to use -- only supports COUNT
     * @throws IllegalArgumentException if what != COUNT
     */

    public StringAggregator(int gbfield, Type gbfieldtype, int afield, Op what) {
        // some code goes here
        if (what!=Op.COUNT) {
            throw new IllegalArgumentException();
        }
        this.gbfield = gbfield;
        this.gbfieldtype = gbfieldtype;
        this.afield = afield;
        this.what = what;
        if (gbfieldtype==null) {
            nogroup = 0;
        }else {
            groups = new HashMap<>();
        }
    }

    /**
     * Merge a new tuple into the aggregate, grouping as indicated in the constructor
     * @param tup the Tuple containing an aggregate field and a group-by field
     */
    public void mergeTupleIntoGroup(Tuple tup) {
        // some code goes here
        StringField string = (StringField) tup.getField(afield);
        if (this.gbfieldtype==null) {
            this.nogroup++;
        }else {
            Field field = tup.getField(gbfield);
            groups.putIfAbsent(field,0);
            groups.put(field,groups.get(field)+1);
        }
    }

    /**
     * Create a OpIterator over group aggregate results.
     *
     * @return a OpIterator whose tuples are the pair (groupVal,
     *   aggregateVal) if using group, or a single (aggregateVal) if no
     *   grouping. The aggregateVal is determined by the type of
     *   aggregate specified in the constructor.
     */
    public OpIterator iterator() {
        // some code goes here
        List<Tuple> tuples = new ArrayList<>();
        TupleDesc td;
        if (this.gbfieldtype==null) {
            td = new TupleDesc(new Type[]{Type.INT_TYPE});
            Tuple a = new Tuple(td);
            a.setField(0,new IntField(nogroup));
            tuples.add(a);
        }else {
            td = new TupleDesc(new Type[]{this.gbfieldtype,Type.INT_TYPE});
            for (Map.Entry<Field,Integer> entry: groups.entrySet()) {
                Tuple a= new Tuple(td);
                if (this.what == Op.COUNT) {
                    a.setField(1,new IntField(entry.getValue()));
                } else {
                    throw new IllegalArgumentException();
                }
                a.setField(0, entry.getKey());
                tuples.add(a);
            }
        }
        return new TupleIterator(td,tuples);
    }

}
